"""
数值方法模块

该模块实现了求解器所需的数值方法，包括：
1. SIMPLE算法
2. PGMRES求解器
3. 预处理方法
4. 收敛判据
"""

import numpy as np
from scipy.sparse.linalg import gmres
from scipy.sparse import csr_matrix
from typing import Tuple, Optional
from .staggered_grid import StaggeredGrid

class NumericalMethods:
    """数值方法类"""
    
    def __init__(self, config: dict):
        """
        初始化数值方法
        
        参数:
            config (dict): 配置参数字典
        """
        # SIMPLE算法参数
        simple_config = config.get('simple', {})
        self.simple_max_iter = simple_config.get('max_iterations', 1000)
        self.simple_tolerance = simple_config.get('tolerance', 1e-6)
        self.alpha_p = simple_config.get('pressure_relaxation', 0.3)
        self.alpha_u = simple_config.get('velocity_relaxation', 0.7)
        
        # PGMRES求解器参数
        pgmres_config = config.get('pgmres', {})
        self.pgmres_max_iter = pgmres_config.get('max_iterations', 1000)
        self.pgmres_tolerance = pgmres_config.get('tolerance', 1e-6)
        self.pgmres_restart = pgmres_config.get('restart', 20)
        
        # 预处理方法参数
        precon_config = config.get('preconditioner', {})
        self.precon_type = precon_config.get('type', 'ilu')
        self.precon_fill = precon_config.get('fill_factor', 1)
        
    def solve_simple(self, grid: StaggeredGrid) -> Tuple[bool, int]:
        """
        使用SIMPLE算法求解流场
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            
        返回:
            Tuple[bool, int]: (是否收敛, 迭代次数)
        """
        # 存储当前时间步的数据
        grid.store_current_step()
        
        for iter in range(self.simple_max_iter):
            # 1. 求解动量方程得到速度场
            self._solve_momentum_equations(grid)
            
            # 2. 求解压力修正方程
            p_correction = self._solve_pressure_correction(grid)
            
            # 3. 修正压力场
            self._correct_pressure(grid, p_correction)
            
            # 4. 修正速度场
            self._correct_velocities(grid, p_correction)
            
            # 5. 检查收敛性
            if self._check_simple_convergence(grid):
                return True, iter + 1
                
        return False, self.simple_max_iter
        
    def solve_pgmres(self, A: csr_matrix, b: np.ndarray) -> Tuple[np.ndarray, int]:
        """
        使用PGMRES求解线性方程组
        
        参数:
            A (csr_matrix): 系数矩阵
            b (np.ndarray): 右端项
            
        返回:
            Tuple[np.ndarray, int]: (解向量, 收敛信息)
        """
        # 构建预处理器
        M = self._build_preconditioner(A)
        
        # 求解方程组
        x, info = gmres(A, b,
                       M=M,
                       maxiter=self.pgmres_max_iter,
                       tol=self.pgmres_tolerance,
                       restart=self.pgmres_restart)
                       
        return x, info
        
    def _solve_momentum_equations(self, grid: StaggeredGrid):
        """求解动量方程"""
        # 获取当前速度场
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        
        # 计算动量方程系数
        A_u, b_u = self._build_axial_momentum_matrix(grid)
        A_v, b_v = self._build_lateral_momentum_matrix(grid)
        
        # 求解动量方程
        u_new, _ = self.solve_pgmres(A_u, b_u)
        v_new, _ = self.solve_pgmres(A_v, b_v)
        
        # 松弛更新
        u = (1 - self.alpha_u) * u + self.alpha_u * u_new
        v = (1 - self.alpha_u) * v + self.alpha_u * v_new
        
        # 更新速度场
        grid.set_axial_velocity_field(u)
        grid.set_lateral_velocity_field(v)
        
    def _solve_pressure_correction(self, grid: StaggeredGrid) -> np.ndarray:
        """求解压力修正方程"""
        # 构建压力修正方程
        A_p, b_p = self._build_pressure_correction_matrix(grid)
        
        # 求解压力修正方程
        p_correction, _ = self.solve_pgmres(A_p, b_p)
        
        return p_correction
        
    def _correct_pressure(self, grid: StaggeredGrid, p_correction: np.ndarray):
        """修正压力场"""
        # 获取当前压力场
        p = grid.get_scalar_field('pressure')
        
        # 松弛更新压力场
        p_new = p + self.alpha_p * p_correction
        
        # 更新压力场
        grid.set_scalar_field('pressure', p_new)
        
    def _correct_velocities(self, grid: StaggeredGrid, p_correction: np.ndarray):
        """修正速度场"""
        # 获取当前速度场
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        
        # 计算速度修正量
        du = self._calculate_axial_velocity_correction(grid, p_correction)
        dv = self._calculate_lateral_velocity_correction(grid, p_correction)
        
        # 更新速度场
        grid.set_axial_velocity_field(u + du)
        grid.set_lateral_velocity_field(v + dv)
        
    def _check_simple_convergence(self, grid: StaggeredGrid) -> bool:
        """检查SIMPLE算法收敛性"""
        # 计算质量守恒残差
        mass_residual = self._calculate_mass_residual(grid)
        
        # 计算动量守恒残差
        momentum_x_residual = self._calculate_axial_momentum_residual(grid)
        momentum_y_residual = self._calculate_lateral_momentum_residual(grid)
        
        # 计算最大残差
        max_residual = max(
            np.max(np.abs(mass_residual)),
            np.max(np.abs(momentum_x_residual)),
            np.max(np.abs(momentum_y_residual))
        )
        
        return max_residual < self.simple_tolerance
        
    def _build_preconditioner(self, A: csr_matrix) -> Optional[csr_matrix]:
        """构建预处理器"""
        if self.precon_type == 'ilu':
            from scipy.sparse.linalg import spilu
            try:
                ilu = spilu(A, fill_factor=self.precon_fill)
                M = lambda x: ilu.solve(x)
                return M
            except:
                print("警告: ILU预处理器构建失败，使用无预处理")
                return None
        elif self.precon_type == 'jacobi':
            # 构建Jacobi预处理器
            diag = A.diagonal()
            M = csr_matrix((1.0/diag, (np.arange(len(diag)), np.arange(len(diag)))))
            return M
        else:
            return None
            
    def _build_axial_momentum_matrix(self, grid: StaggeredGrid) -> Tuple[csr_matrix, np.ndarray]:
        """
        构建轴向动量方程的系数矩阵
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            
        返回:
            Tuple[csr_matrix, np.ndarray]: (系数矩阵, 右端项)
        """
        # 获取网格尺寸
        nx, ny = grid.get_dimensions()
        n = (nx+1) * ny  # 轴向速度网格点数
        
        # 初始化稀疏矩阵数据
        data = []
        row_ind = []
        col_ind = []
        b = np.zeros(n)
        
        # 获取物理量
        alpha = grid.get_scalar_field('void_fraction')
        rho = grid.get_scalar_field('density')
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        P = grid.get_scalar_field('pressure')
        
        # 遍历内部点
        for i in range(nx - 1):
            for j in range(ny):
                idx = i * ny + j
                
                # 时间项系数
                dt = grid.dt
                data.append(alpha[i,j] * rho[i,j] / dt)
                row_ind.append(idx)
                col_ind.append(idx)
                
                # 对流项系数 (upwind scheme)
                if u[i,j] > 0:
                    data.append(-alpha[i,j] * rho[i,j] * u[i,j] / grid.dx)
                    row_ind.append(idx)
                    col_ind.append(idx-ny)
                else:
                    data.append(alpha[i,j] * rho[i,j] * u[i,j] / grid.dx)
                    row_ind.append(idx)
                    col_ind.append(idx+ny)
                    
                # 压力梯度项
                b[idx] -= alpha[i,j] * (P[i+1,j] - P[i,j]) / grid.dx
                
                # 重力项
                g = 9.81
                b[idx] -= alpha[i,j] * rho[i,j] * g
                
                # 壁面摩擦阻力项
                F_w = grid.get_source_term('wall_friction_x')[i,j]
                b[idx] -= F_w
                
                # 界面剪切力项
                F_i = grid.get_source_term('interfacial_force_x')[i,j]
                b[idx] += F_i
                
                # 相变动量项
                Gamma = grid.get_source_term('phase_change_momentum_x')[i,j]
                b[idx] -= Gamma
                
        # 构建稀疏矩阵
        A = csr_matrix((data, (row_ind, col_ind)), shape=(n, n))
        
        return A, b
        
    def _build_lateral_momentum_matrix(self, grid: StaggeredGrid) -> Tuple[csr_matrix, np.ndarray]:
        """
        构建横向动量方程的系数矩阵
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            
        返回:
            Tuple[csr_matrix, np.ndarray]: (系数矩阵, 右端项)
        """
        # 获取网格尺寸
        nx, ny = grid.get_dimensions()
        n = nx * (ny + 1)  # 横向速度网格点数
        
        # 初始化稀疏矩阵数据
        data = []
        row_ind = []
        col_ind = []
        b = np.zeros(n)
        
        # 获取物理量
        alpha = grid.get_scalar_field('void_fraction')
        rho = grid.get_scalar_field('density')
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        P = grid.get_scalar_field('pressure')
        
        # 遍历内部点
        for i in range(nx):
            for j in range(1, ny):
                idx = i * (ny + 1) + j
                
                # 时间项系数
                dt = grid.dt
                data.append(alpha[i, j] * rho[i, j] / dt)
                row_ind.append(idx)
                col_ind.append(idx)
                
                # 对流项系数 (简化处理，仅考虑主要项)
                if u[i, j] > 0:
                    if j > 0:  # 确保不超出边界
                        data.append(-alpha[i, j] * rho[i, j] * u[i, j] / grid.dx)
                        row_ind.append(idx)
                        col_ind.append(idx - (ny + 1))  # 确保 idx - (ny + 1) 不小于 0
                else:
                    if j < ny:  # 确保不超出边界
                        data.append(alpha[i, j] * rho[i, j] * u[i, j] / grid.dx)
                        row_ind.append(idx)
                        col_ind.append(idx + (ny + 1))  # 确保 idx + (ny + 1) 不超过 n
                    
                # 压力梯度项
                if j < ny - 1:
                    b[idx] -= alpha[i, j] * (P[i, j + 1] - P[i, j]) / grid.dy
                
                # 壁面摩擦阻力项
                F_w = grid.get_source_term('wall_friction_y')[i, j]
                b[idx] -= F_w
                
                # 界面剪切力项
                F_i = grid.get_source_term('interfacial_force_y')[i, j]
                b[idx] += F_i
                
                # 相变动量项
                Gamma = grid.get_source_term('phase_change_momentum_y')[i, j]
                b[idx] -= Gamma
                
                # 湍流混合动量项
                T_m = grid.get_source_term('turbulent_mixing_momentum_y')[i, j]
                b[idx] -= T_m
                
        # 构建稀疏矩阵
        A = csr_matrix((data, (row_ind, col_ind)), shape=(n, n))
        
        return A, b
        
    def _build_pressure_correction_matrix(self, grid: StaggeredGrid) -> Tuple[csr_matrix, np.ndarray]:
        """
        构建压力修正方程的系数矩阵
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            
        返回:
            Tuple[csr_matrix, np.ndarray]: (系数矩阵, 右端项)
        """
        # 获取网格尺寸
        nx, ny = grid.get_dimensions()
        n = nx * ny  # 压力网格点数
        
        # 初始化稀疏矩阵数据
        data = []
        row_ind = []
        col_ind = []
        b = np.zeros(n)
        
        # 获取物理量
        alpha = grid.get_scalar_field('void_fraction')
        rho = grid.get_scalar_field('density')
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        
        # 遍历内部点
        for i in range(nx):
            for j in range(ny):
                idx = i * ny + j
                
                # 中心点系数
                a_p = 0.0
                
                # 轴向速度贡献
                if i > 0:
                    d_w = alpha[i-1,j] / (grid.dx * grid.get_source_term('axial_momentum_coeff')[i-1,j])
                    a_p += d_w
                    data.append(-d_w)
                    row_ind.append(idx)
                    col_ind.append(idx-ny)
                    
                if i < nx-1:
                    d_e = alpha[i+1,j] / (grid.dx * grid.get_source_term('axial_momentum_coeff')[i+1,j])
                    a_p += d_e
                    data.append(-d_e)
                    row_ind.append(idx)
                    col_ind.append(idx+ny)
                    
                # 横向速度贡献
                if j > 0:
                    d_s = alpha[i,j-1] / (grid.dy * grid.get_source_term('lateral_momentum_coeff')[i,j-1])
                    a_p += d_s
                    data.append(-d_s)
                    row_ind.append(idx)
                    col_ind.append(idx-1)
                    
                if j < ny-1:
                    d_n = alpha[i,j+1] / (grid.dy * grid.get_source_term('lateral_momentum_coeff')[i,j+1])
                    a_p += d_n
                    data.append(-d_n)
                    row_ind.append(idx)
                    col_ind.append(idx+1)
                    
                # 中心点系数
                data.append(a_p)
                row_ind.append(idx)
                col_ind.append(idx)
                
                # 右端项 (质量不平衡)
                b[idx] = -(
                    (alpha[i+1,j]*u[i+1,j] - alpha[i,j]*u[i,j]) / grid.dx +
                    (alpha[i,j+1]*v[i,j+1] - alpha[i,j]*v[i,j]) / grid.dy
                )
                
        # 构建稀疏矩阵
        A = csr_matrix((data, (row_ind, col_ind)), shape=(n, n))
        
        return A, b
        
    def _calculate_velocity_corrections(self, grid: StaggeredGrid, p_corr: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        计算速度修正值
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            p_corr (np.ndarray): 压力修正值
            
        返回:
            Tuple[np.ndarray, np.ndarray]: (轴向速度修正, 横向速度修正)
        """
        # 获取网格尺寸
        nx, ny = grid.get_dimensions()
        
        # 初始化速度修正数组
        u_corr = np.zeros((nx+1, ny))
        v_corr = np.zeros((nx, ny+1))
        
        # 计算轴向速度修正
        for i in range(1, nx):
            for j in range(ny):
                d_u = grid.get_source_term('axial_momentum_coeff')[i,j]
                p_grad = (p_corr[i*ny + j] - p_corr[(i-1)*ny + j]) / grid.dx
                u_corr[i,j] = -d_u * p_grad
                
        # 计算横向速度修正
        for i in range(nx):
            for j in range(1, ny):
                d_v = grid.get_source_term('lateral_momentum_coeff')[i,j]
                p_grad = (p_corr[i*ny + j] - p_corr[i*ny + j-1]) / grid.dy
                v_corr[i,j] = -d_v * p_grad
                
        return u_corr, v_corr
        
    def _calculate_axial_velocity_correction(self, grid: StaggeredGrid, p_correction: np.ndarray) -> np.ndarray:
        """计算轴向速度修正量"""
        # TODO: 实现轴向速度修正量的计算
        pass
        
    def _calculate_lateral_velocity_correction(self, grid: StaggeredGrid, p_correction: np.ndarray) -> np.ndarray:
        """计算横向速度修正量"""
        # TODO: 实现横向速度修正量的计算
        pass
        
    def _calculate_mass_residual(self, grid: StaggeredGrid) -> np.ndarray:
        """计算质量守恒残差"""
        # TODO: 实现质量守恒残差的计算
        pass
        
    def _calculate_axial_momentum_residual(self, grid: StaggeredGrid) -> np.ndarray:
        """计算轴向动量守恒残差"""
        # TODO: 实现轴向动量守恒残差的计算
        pass
        
    def _calculate_lateral_momentum_residual(self, grid: StaggeredGrid) -> np.ndarray:
        """计算横向动量守恒残差"""
        # TODO: 实现横向动量守恒残差的计算
        pass
        
    def _calculate_residuals(self, grid: StaggeredGrid) -> dict:
        """
        计算各守恒方程的残差
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            
        返回:
            dict: 包含各守恒方程残差的字典
        """
        # 获取网格尺寸
        nx, ny = grid.get_dimensions()
        
        # 初始化残差字典
        residuals = {
            'mass': 0.0,
            'axial_momentum': 0.0,
            'lateral_momentum': 0.0,
            'energy': 0.0
        }
        
        # 获取物理量
        alpha = grid.get_scalar_field('void_fraction')
        rho = grid.get_scalar_field('density')
        u = grid.get_axial_velocity_field()
        v = grid.get_lateral_velocity_field()
        P = grid.get_scalar_field('pressure')
        h = grid.get_scalar_field('enthalpy')
        
        # 计算质量守恒残差
        mass_res = 0.0
        for i in range(nx):
            for j in range(ny):
                mass_flux_x = (alpha[i+1,j]*rho[i+1,j]*u[i+1,j] - alpha[i,j]*rho[i,j]*u[i,j]) / grid.dx
                mass_flux_y = (alpha[i,j+1]*rho[i,j+1]*v[i,j+1] - alpha[i,j]*rho[i,j]*v[i,j]) / grid.dy
                mass_res += (mass_flux_x + mass_flux_y)**2
        residuals['mass'] = np.sqrt(mass_res / (nx * ny))
        
        # 计算轴向动量守恒残差
        axial_res = 0.0
        for i in range(1, nx):
            for j in range(ny):
                # 动量通量
                momentum_flux = alpha[i,j] * rho[i,j] * u[i,j] * (u[i+1,j] - u[i-1,j]) / (2*grid.dx)
                
                # 压力梯度
                pressure_grad = alpha[i,j] * (P[i+1,j] - P[i,j]) / grid.dx
                
                # 重力项
                gravity = alpha[i,j] * rho[i,j] * 9.81
                
                # 壁面摩擦力
                F_w = grid.get_source_term('wall_friction_x')[i,j]
                
                # 界面剪切力
                F_i = grid.get_source_term('interfacial_force_x')[i,j]
                
                # 相变动量
                Gamma = grid.get_source_term('phase_change_momentum_x')[i,j]
                
                # 计算残差
                res = momentum_flux + pressure_grad + gravity + F_w - F_i + Gamma
                axial_res += res**2
                
        residuals['axial_momentum'] = np.sqrt(axial_res / ((nx-1) * ny))
        
        # 计算横向动量守恒残差
        lateral_res = 0.0
        for i in range(nx):
            for j in range(1, ny):
                # 动量通量
                momentum_flux = alpha[i,j] * rho[i,j] * v[i,j] * (v[i,j+1] - v[i,j-1]) / (2*grid.dy)
                
                # 压力梯度
                pressure_grad = alpha[i,j] * (P[i,j+1] - P[i,j]) / grid.dy
                
                # 壁面摩擦力
                F_w = grid.get_source_term('wall_friction_y')[i,j]
                
                # 界面剪切力
                F_i = grid.get_source_term('interfacial_force_y')[i,j]
                
                # 相变动量
                Gamma = grid.get_source_term('phase_change_momentum_y')[i,j]
                
                # 湍流混合动量
                T_m = grid.get_source_term('turbulent_mixing_momentum_y')[i,j]
                
                # 计算残差
                res = momentum_flux + pressure_grad + F_w - F_i + Gamma + T_m
                lateral_res += res**2
                
        residuals['lateral_momentum'] = np.sqrt(lateral_res / (nx * (ny-1)))
        
        # 计算能量守恒残差
        energy_res = 0.0
        for i in range(nx):
            for j in range(ny):
                # 能量通量
                energy_flux_x = alpha[i,j] * rho[i,j] * u[i,j] * h[i,j] * (h[i+1,j] - h[i-1,j]) / (2*grid.dx)
                energy_flux_y = alpha[i,j] * rho[i,j] * v[i,j] * h[i,j] * (h[i,j+1] - h[i,j-1]) / (2*grid.dy)
                
                # 导热
                q_cond = grid.get_source_term('heat_conduction')[i,j]
                
                # 壁面传热
                q_w = grid.get_source_term('wall_heat_transfer')[i,j]
                
                # 界面传热
                q_i = grid.get_source_term('interfacial_heat_transfer')[i,j]
                
                # 相变潜热
                q_phase = grid.get_source_term('phase_change_heat')[i,j]
                
                # 计算残差
                res = energy_flux_x + energy_flux_y - q_cond - q_w + q_i - q_phase
                energy_res += res**2
                
        residuals['energy'] = np.sqrt(energy_res / (nx * ny))
        
        return residuals